package proxy

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"strconv"
	"strings"
	"time"

	"gitee.com/baixudong/gospider/kinds"
	"gitee.com/baixudong/gospider/re"
	"gitee.com/baixudong/gospider/thread"
	"gitee.com/baixudong/gospider/tools"
	netProxy "golang.org/x/net/proxy"
)

type Type = int

var (
	Socks5 Type
	Http   Type = 1
)

type ClientOption struct {
	Usr         string      //用户名
	Pwd         string      //密码
	IpWhite     []net.IP    //白名单 192.168.1.1,192.168.1.2
	PoolNum     int         //并发数量
	PoolTimeout int         //pool wait task time
	Dialer      *net.Dialer //连接的Dialer
	LocalAddr   string      //本地网卡出口
	Port        int         //代理端口
	Host        string      //代理host
}
type Client struct {
	Proxy    string                 //代理ip 192.168.1.50:8888
	GetProxy func() (string, error) //代理ip 116.62.55.139:8888
	Debug    bool                   //是否打印debug
	Type     Type                   //协议类型
	Dialer   *net.Dialer            //连接的Dialer

	listener net.Listener //Listener 服务

	Err         error //错误
	basic       []byte
	usr         string
	pwd         string
	verify      bool
	ipWhite     *kinds.Set[string]
	poolNum     int
	poolTimeout int
	ctx         context.Context
	cnl         context.CancelFunc
}

func NewClient(pre_ctx context.Context, options ...ClientOption) (*Client, error) {
	var option ClientOption
	if len(options) > 0 {
		option = options[0]
	}
	if pre_ctx == nil {
		pre_ctx = context.TODO()
	}
	ctx, cnl := context.WithCancel(pre_ctx)
	server := Client{}
	server.ctx = ctx
	server.cnl = cnl
	if option.Usr != "" && option.Pwd != "" {
		server.basic = []byte(base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%s:%s", option.Usr, option.Pwd))))
		server.usr = option.Usr
		server.pwd = option.Pwd
		server.verify = true
	}
	if option.PoolNum == 0 {
		server.poolNum = 1000000
	} else {
		server.poolNum = option.PoolNum
	}
	if option.PoolTimeout == 0 {
		server.poolTimeout = 60
	} else {
		server.poolTimeout = option.PoolTimeout
	}
	if len(option.IpWhite) > 0 {
		ip_whites := kinds.NewSet[string]()
		for _, ip_white := range option.IpWhite {
			ip_whites.Add(ip_white.String())
		}
		server.ipWhite = ip_whites
	}

	if option.Dialer == nil {
		server.Dialer = &net.Dialer{
			Timeout:   time.Duration(8) * time.Second,
			KeepAlive: time.Duration(10) * time.Second,
		}
	} else {
		server.Dialer = option.Dialer
	}

	if option.LocalAddr != "" {
		if !strings.Contains(option.LocalAddr, ":") {
			option.LocalAddr += ":0"
		}
		localaddr, err := net.ResolveTCPAddr("tcp", option.LocalAddr)
		if err != nil {
			return nil, err
		}
		server.Dialer.LocalAddr = localaddr
	}

	l, err := net.Listen("tcp", fmt.Sprintf("%s:%d", option.Host, option.Port)) //监听本地端口
	if err != nil {
		return nil, err
	}
	server.listener = l
	return &server, nil
}

// 代理监听的端口
func (obj *Client) Addr() string {
	return obj.listener.Addr().String()
}

func (obj *Client) Run() error {
	defer obj.Close()
	pool := thread.NewClient(obj.ctx, obj.poolNum, thread.ClientOption{Timeout: obj.poolTimeout})
	pool.Debug = obj.Debug
	for {
		select {
		case <-obj.ctx.Done():
			obj.Err = obj.ctx.Err()
			return obj.Err
		default:
			client, err := obj.listener.Accept() //接受数据
			if err != nil {
				obj.Err = err
				return err
			}
			switch obj.Type {
			case Http:
				_, err = pool.Write(&thread.Task{
					Func: obj.httpHandle,
					Args: []any{client},
				})
				if err != nil {
					obj.Err = err
					return obj.Err
				}
			case Socks5:
				_, err := pool.Write(&thread.Task{
					Func: obj.sockes5Handle,
					Args: []any{client},
				})
				if err != nil {
					obj.Err = err
					return obj.Err
				}
			default:
				obj.Err = errors.New("type error")
				return obj.Err
			}
		}
	}
}
func (obj *Client) Close() {
	obj.listener.Close()
	obj.cnl()
}
func (obj *Client) Done() <-chan struct{} {
	return obj.ctx.Done()
}

func (obj *Client) whiteVerify(client net.Conn) bool {
	if obj.ipWhite == nil {
		return true
	}
	host, _, err := net.SplitHostPort(client.RemoteAddr().String())
	if err != nil || !obj.ipWhite.Has(host) {
		return false
	}
	return true
}
func (obj *Client) readAll(client net.Conn) (bytes.Buffer, error) {
	var readCon bytes.Buffer
	readNum := 1024
	cc_content := make([]byte, readNum)
	for {
		n, err := client.Read(cc_content)
		if err != nil {
			return readCon, err
		}
		readCon.Write(cc_content[:n])
		if bytes.Contains(readCon.Bytes(), []byte("\r\n\r\n")) || n < readNum {
			break
		}
	}
	return readCon, nil
}
func (obj *Client) parseFirstLine(readCon bytes.Buffer) ([][]byte, []byte, error) {
	read_index := bytes.Index(readCon.Bytes(), []byte("\r\n"))
	if read_index == -1 || read_index == 0 {
		return nil, nil, errors.New("not found read_index")
	}
	read_nn := readCon.Bytes()[:read_index]
	read_lls := bytes.Split(read_nn, []byte(" "))
	return read_lls, read_nn, nil
}

// 返回:请求所有内容,第一行的内容被" "分割的数组,第一行的内容,error
func (obj *Client) verifyPwd(client net.Conn, readCon bytes.Buffer, read_lls [][]byte) error {
	if !bytes.Contains(readCon.Bytes(), obj.basic) { //验证密码是否正确
		client.Write([]byte(fmt.Sprintf("%s 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic\r\n\r\n", read_lls[2])))
		return errors.New("auth verify fail")
	}
	return nil
}
func (obj *Client) parseUrl(read_lls [][]byte) (*url.URL, string, error) {
	hostPortURL, err := url.Parse(tools.BytesToString(read_lls[1])) //解析url
	if err != nil {
		return hostPortURL, "", err
	}
	var ip_addr string
	if bytes.Contains(read_lls[0], []byte("CONNECT")) {
		ip_addr = hostPortURL.Scheme + ":" + hostPortURL.Opaque
	} else if !strings.Contains(hostPortURL.Host, ":") { //host不带端口， 默认80
		ip_addr = hostPortURL.Host + ":80"
	} else {
		ip_addr = hostPortURL.Host
	}
	return hostPortURL, ip_addr, nil
}

func (obj *Client) getHttpWriteData(readCon bytes.Buffer, hostPortURL *url.URL, read_lls [][]byte, read_nn []byte) []byte {
	_, afterData, ok := bytes.Cut(read_lls[1], []byte(hostPortURL.Host))
	if ok {
		read_lls[1] = afterData
	} else {
		read_lls[1] = []byte(hostPortURL.Path)
	}
	tempData := bytes.Replace(readCon.Bytes(), read_nn, bytes.Join(read_lls, []byte(" ")), 1)
	tempData = tools.StringToBytes(re.Sub("\r\nProxy-.*?\r\n", "\r\n", tools.BytesToString(tempData)))
	return tempData
}
func (obj *Client) getProxyWriteData(readCon bytes.Buffer, pwd string) []byte {
	return tools.StringToBytes(strings.Replace(re.Sub("\r\nProxy-.*?\r\n", "\r\n", readCon.String()), "\r\n\r\n", fmt.Sprintf("\r\nProxy-Authorization: Basic %s\r\n\r\n", pwd), 1))
}

func (obj *Client) getHttpProxyConn(ctx context.Context, ipUrl *url.URL) (net.Conn, error) {
	return obj.Dialer.DialContext(ctx, "tcp", net.JoinHostPort(ipUrl.Hostname(), ipUrl.Port()))
}
func (obj *Client) getSocksProxyConn(ctx context.Context, ipUrl *url.URL, serverAddr string) (net.Conn, error) {
	var auth *netProxy.Auth
	if ipUrl.User != nil {
		auth = new(netProxy.Auth)
		auth.User = ipUrl.User.Username()
		if p, ok := ipUrl.User.Password(); ok {
			auth.Password = p
		}
	}
	dial, err := netProxy.SOCKS5("tcp", net.JoinHostPort(ipUrl.Hostname(), ipUrl.Port()), auth, obj.Dialer)
	if err != nil {
		return nil, err
	}
	return dial.Dial("tcp", serverAddr)
}
func (obj *Client) httpHandle(ctx context.Context, client net.Conn) {
	if client == nil {
		return
	}
	defer client.Close()
	if !obj.whiteVerify(client) {
		return
	}
	var err error
	var ip_addr string
	if obj.GetProxy != nil {
		ip_addr, err = obj.GetProxy()
		if err != nil {
			return
		}
	} else if obj.Proxy != "" {
		ip_addr = obj.Proxy
	}
	var server net.Conn
	if ip_addr == "" { //使用本地转发的逻辑
		readCon, err := obj.readAll(client)
		if err != nil {
			return
		}
		read_lls, read_nn, err := obj.parseFirstLine(readCon)
		if err != nil {
			return
		}
		if obj.verify {
			err = obj.verifyPwd(client, readCon, read_lls)
			if err != nil {
				return
			}
		}
		hostPortURL, ip_addr, err := obj.parseUrl(read_lls)
		if err != nil {
			return
		}
		server, err = obj.Dialer.DialContext(ctx, "tcp", ip_addr)
		if err != nil {
			return
		}
		defer server.Close()
		if bytes.Contains(read_lls[0], []byte("CONNECT")) {
			_, err = client.Write([]byte(fmt.Sprintf("%s 200 Connection established\r\n\r\n", read_lls[2])))
			if err != nil {
				return
			}
		} else {
			_, err := server.Write(obj.getHttpWriteData(readCon, hostPortURL, read_lls, read_nn))
			if err != nil {
				return
			}
		}
	} else { //使用代理转发的逻辑
		var readCon bytes.Buffer
		var err error
		var ipUrl *url.URL
		if obj.verify {
			if readCon, err = obj.readAll(client); err != nil {
				return
			}
			read_lis, _, err := obj.parseFirstLine(readCon)
			if err != nil {
				return
			}
			if err = obj.verifyPwd(client, readCon, read_lis); err != nil {
				return
			}
		}
		if ipUrl, err = url.Parse(ip_addr); err != nil {
			return
		}
		if server, err = obj.getHttpProxyConn(ctx, ipUrl); err != nil {
			return
		}
		defer server.Close()
		var pwd string
		if ipUrl.User != nil {
			if _, ok := ipUrl.User.Password(); ok {
				pwd = tools.Base64Encode(ipUrl.User.String())
			}
		}
		if pwd != "" && obj.verify {
			if _, err = server.Write(obj.getProxyWriteData(readCon, pwd)); err != nil {
				return
			}
		} else if obj.verify {
			if _, err = server.Write(readCon.Bytes()); err != nil {
				return
			}
		} else if pwd != "" {
			if readCon, err = obj.readAll(client); err != nil {
				return
			}
			if _, err = server.Write(obj.getProxyWriteData(readCon, pwd)); err != nil {
				return
			}
		}
	}
	go func() { //服务端到客户端
		defer server.Close()
		defer client.Close()
		io.Copy(client, server)
	}()
	io.Copy(server, client) //客户端发送服务端
}

func (obj *Client) getSocketAddr(client net.Conn) (string, error) {
	reaCon, err := obj.readAll(client)
	if err != nil {
		return "", err
	}
	byteCon := reaCon.Bytes()
	byteLen := reaCon.Len()
	var host string
	switch byteCon[3] {
	case 1: //IP V4
		host = net.IPv4(byteCon[4], byteCon[5], byteCon[6], byteCon[7]).String()
	case 3: //Domain
		host = string(byteCon[5 : byteLen-2]) //byteCon[4]表示域名的长度
	case 4: //IP V6
		host = net.IP{byteCon[4], byteCon[5], byteCon[6], byteCon[7], byteCon[8], byteCon[9], byteCon[10], byteCon[11], byteCon[12], byteCon[13], byteCon[14], byteCon[15], byteCon[16], byteCon[17], byteCon[18], byteCon[19]}.String()
	default:
		return "", errors.New("socket未知的协议")
	}
	return net.JoinHostPort(host, strconv.Itoa(int(byteCon[byteLen-2])<<8|int(byteCon[byteLen-1]))), nil
}
func (obj *Client) verifySocket(client net.Conn) error {
	var err error
	reaCon, err := obj.readAll(client)
	if err != nil {
		return err
	}
	byteCon := reaCon.Bytes()
	if int(byteCon[0]) != 5 {
		return errors.New("不是socket5 协议")
	}
	if obj.verify { //验证用户名密码
		_, err = client.Write([]byte{byteCon[0], 2})
		if err != nil {
			return err
		}
		reaCon, err = obj.readAll(client)
		if err != nil {
			return err
		}
		byteCon = reaCon.Bytes()
		userLen := byteCon[1]
		user := byteCon[2 : userLen+2]
		pass := byteCon[userLen+3 : userLen+3+byteCon[userLen+2]]
		if tools.BytesToString(user) != obj.usr || tools.BytesToString(pass) != obj.pwd {
			client.Write([]byte{byteCon[0], 0xff}) //用户名密码错误
			return errors.New("用户名密码错误")
		}
	}
	_, err = client.Write([]byte{byteCon[0], 0}) //协商成功
	return err
}
func (obj *Client) sockes5Handle(ctx context.Context, client net.Conn) {
	if client == nil {
		return
	}
	defer client.Close()
	if !obj.whiteVerify(client) {
		return
	}
	var ip_addr string
	var err error
	if obj.GetProxy != nil {
		if ip_addr, err = obj.GetProxy(); err != nil {
			return
		}
	} else if obj.Proxy != "" {
		ip_addr = obj.Proxy
	}
	var server net.Conn
	if err = obj.verifySocket(client); err != nil {
		return
	}
	if ip_addr == "" {
		serverAddr, err := obj.getSocketAddr(client)
		if err != nil {
			return
		}
		server, err = obj.Dialer.DialContext(ctx, "tcp", serverAddr)
		if err != nil {
			return
		}
		defer server.Close()
		_, err = client.Write([]byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}) //响应客户端连接成功
		if err != nil {
			return
		}
	} else {
		serverAddr, err := obj.getSocketAddr(client)
		if err != nil {
			return
		}
		_, err = client.Write([]byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}) //响应客户端连接成功
		if err != nil {
			return
		}
		ipUrl, err := url.Parse(ip_addr)
		if err != nil {
			return
		}
		if server, err = obj.getSocksProxyConn(ctx, ipUrl, serverAddr); err != nil {
			return
		}
		defer server.Close()
	}
	go func() { //服务端到客户端
		defer server.Close()
		defer client.Close()
		io.Copy(client, server)
	}()
	io.Copy(server, client) //客户端发送服务端
}
